#include "predict.h"
#include "common/utils.h"
#include <stdexcept>
#include <iostream>
#include <sstream>
#include <unistd.h>


namespace parser {
// Predict Method Implementation

void PredictParser::ExtractCommonFactors() 
{
    // If a non-terminal has multiple candidate generation expressions with a
    // same common prefix/factor then extract common factor and change syntax rules.

}

void PredictParser::RemoveLeftRecursion()
{
    // TODO: remove left recursion.
}


void PredictParser::GetFirstSet(char symbol, std::unordered_set<char> &first_set) 
{
    /**
     * first set is a set of symbols that may appear at the first position of
     * all expressions generated by a non-terminal in the syntax.
     */
    if (syntax.vts.count(symbol) || symbol == syntax.epsilon)
    {
        // first set of a terminal symbol is itself.
        first_set.insert(symbol);
        return;
    }
    char cur_vn = symbol;
    for (const auto &rule : syntax.rules)
    {
        if (cur_vn != rule.first) continue;
        auto expr = rule.second;
        char first_symbol = expr.front();
        if (syntax.vts.count(first_symbol)
            || first_symbol == syntax.epsilon && expr.size() == 1)
        {
            /*
                for A -> αβ :
                    if α is terminal symbol
                    then α ∈ FIRST(A)
                for A -> ε there is: ε ∈ FIRST(A)
            */
            first_set.insert(first_symbol);
            continue;
        }
        if (syntax.vns.count(first_symbol))
        {
            // If the first_symbol is a non-terminal symbol and FIRST(first_symbol)
            // contains empty words ε then we need to look at the next symbol.
            bool has_empty_symbol = true;
            int idx = 0;
            /*
                for A -> αβ :
                    if α is non-terminal and ε ∈ FIRST(α)
                    then FIRST(α) \ {ε} ⊂ FIRST(A) and FIRST(β) ⊂ FIRST(A)
                for A -> α1 α2 α3 ... αn:
                    loop i from 1 to n:
                        if αi is non-terminal and ε ∈ FIRST(αi)
                        then FIRST(αi) \ {ε} ⊂ FIRST(A), i += 1
                        if αi is terminal then αi ∈ FIRST(αi) and break
                    if i >= n then ε ∈ FIRST(A)  because all the symbols in expr may be ε
            */
            while (has_empty_symbol && idx < expr.size())
            {
                std::unordered_set<char> next_first_set;
                first_symbol = expr[idx++];
                // recursive call this function to obtain the first set of non-terminal `first_symbol`
                // so we need to ensure that there's no left recursion in syntax.
                GetFirstSet(first_symbol, next_first_set);
                if (next_first_set.count(syntax.epsilon))
                {
                    has_empty_symbol = true;
                    if (idx < expr.size() - 1)
                    {
                        next_first_set.erase(syntax.epsilon);
                    }
                }
                else
                {
                    has_empty_symbol = false;
                }
                for (auto nef : next_first_set)
                {
                    first_set.insert(nef);
                }
            }
            continue;
        }

        throw std::runtime_error(
            "Syntax Rules Error: There is a symbol that isn't in rules, please check."
        );
    }
}


void PredictParser::GetFollowSet(char symbol, std::unordered_set<char> &follow_set) 
{
    /**
     * follow set of X is a set of terminal that may appear after X
     */
    // a terminal symbol doesn't have FOLLOW set.
    if (syntax.vts.count(symbol)) return;
    if (symbol == syntax.start)
    {
        follow_set.insert('#');
    }
    for (const auto &rule : syntax.rules)
    {
        auto expr = rule.second;
        auto idx = expr.find_first_of(symbol);
        if (idx == std::string::npos) continue;

        auto left_vn = rule.first;
        while (idx < expr.size() - 1 && idx != std::string::npos)
        {
            /*
                for A -> α B β ... :
                    FIRST(β) \ {ε} ⊂ FOLLOW(B)

                for A -> α B β1 β2..., and all of {β1,β2,...} has epsilon:
                    FOLLOW(A) ⊂ FOLLOW(B)

                that means,
                    while true do
                         FIRST(β) \ {ε} ⊂ FOLLOW(B)  // where β is symbol next to B
                         if ε not ∈ FIRST(β) then break;
                         if β is last then FOLLOW(A) ⊂ FOLLOW(B)
                         β = get_next_symbol_of_β()
                    done

                for A -> α B:
                    FOLLOW(A) ⊂ FOLLOW(B)
            */

            size_t next_idx = idx + 1;
            char next_symbol = expr[next_idx];

            bool first_has_empty = true;
            std::unordered_set<char> tmp;
            while (next_idx < expr.size() && first_has_empty)
            {
                if (syntax.vts.count(next_symbol))
                {
                    first_has_empty = false;
                }
                else
                {
                    first_has_empty = symbol_sets[next_symbol].first_has_empty;
                }
                tmp.insert(next_symbol);
                next_symbol = expr[++next_idx];
            }

            for (char x : tmp)
            {
                // collect all FIRST() \ {ε} to FOLLOW()
                for (char c : symbol_sets[x].first_set)
                {
                    if (c != syntax.epsilon)
                    {
                        follow_set.insert(c);
                    }
                }
            }
            if (next_idx >= expr.size())
            {
                // if last symbol has epsilon,
                // then collect left-symbol's FOLLOW to this FOLLOW
                for (char c : symbol_sets[left_vn].follow_set)
                {
                    follow_set.insert(c);
                }
            }
            idx = expr.find_first_of(symbol, next_idx);
        }
        if (idx >= expr.size() - 1 && idx != std::string::npos)
        {
            // if last symbol is B,
            // then collect left-symbol(A)'s FOLLOW to B's FOLLOW
            for (char c : symbol_sets[left_vn].follow_set)
            {
                follow_set.insert(c);
            }
        }
    }
}


void PredictParser::MakePredictTable()
{
    /**
     * directly make predict table without SELECT sets.
     * non-terminal symbol as row, terminal symbol as column.
     */
    auto get_first_set_formula = [&] (std::string &form, std::unordered_set<char> &res)
    {
        /**
         * this function is used to get FIRST(α) of rule A -> α
         */
        for (int i = 0; i < form.size(); ++i)
        {
            char c = form[i];
            auto &tmp = symbol_sets[c].first_set;
            for (char x : tmp)
            {
                if (x == syntax.epsilon) continue;
                res.insert(x);
            }
            if (!symbol_sets[c].first_has_empty) 
            {
                break;
            }
            if (i >= form.size() - 1) 
            {
                res.insert(syntax.epsilon);
            }
        }
    };

    for (int index = 0; index < syntax.rules.size(); ++index)
    {
        auto &rule = syntax.rules[index];
        /*
         * for each rule A -> α, we do:
         *   for each terminal a ∈ FIRST(α), add rule A -> α to table position M[A,a]
         *   if ε ∈ FIRST(α), then:
         *     for each symbol b ∈ FOLLOW(A), add rule A -> α to table position M[A,b]
         */
        std::unordered_set<char> tmp_first;
        auto formula = rule.second;
        auto left = rule.first; // here left is 'A'
        get_first_set_formula(formula, tmp_first);
        // we assume that the index of each rule in syntax.rules serves as its identifier.
        for (char a : tmp_first)
        {
            if (a == syntax.epsilon) continue;
            table.emplace(common::CharsToUint16(left, a), index);
        }
        if (tmp_first.count(syntax.epsilon))
        {
            for (char b : symbol_sets[left].follow_set)
            {
                table.emplace(common::CharsToUint16(left, b), index);
            }
        }
    }
}


void PredictParser::Parse(std::string &expression)
{
    // get first sets.
    for (char vn : syntax.vns)
    {
        GetFirstSet(vn, symbol_sets[vn].first_set);
        symbol_sets[vn].first_has_empty = symbol_sets[vn].first_set.count(syntax.epsilon);
    }
    for (char vt : syntax.vts)
    {
        GetFirstSet(vt, symbol_sets[vt].first_set);
        symbol_sets[vt].first_has_empty = false;
    }
    GetFirstSet(syntax.epsilon, symbol_sets[syntax.epsilon].first_set);
    symbol_sets[syntax.epsilon].first_has_empty = true;
    // get follow sets.
    for (char vn : syntax.vns)
    {
        GetFollowSet(vn, symbol_sets[vn].follow_set);
    }

    // construct predict-table.
    MakePredictTable();

    // init character stack
    std::vector<char> stk;
    stk.push_back('#');
    stk.push_back(syntax.start);

    // make an expression copy to buffer.
    std::string expr_buffer(expression);
    expr_buffer.push_back('#');

    int step = 0, idx = 0, success = 0;
    auto print_process = [&] (int rule_idx)
    {
        /**
         * A function to print the process of analyzing.
         */
        const int COL_WIDTHS[3] = {20, 16, 16};
        using namespace common;

        if (step == 0)
        {
            std::cout << PadColumn("Character Stack", COL_WIDTHS[0])
                      << PadColumn("Buffer", COL_WIDTHS[1])
                      << PadColumn("Rule Used", COL_WIDTHS[2]) << "\n";
        }
        std::stringstream cs;
        for (char x : stk) cs << x;
        std::string _formula("-");
        if (rule_idx >= 0)
        {
            std::stringstream _;
            _ << "[" << syntax.rules[rule_idx].first << " -> " << syntax.rules[rule_idx].second << "]";
            _formula = _.str();
        }
        std::cout << PadColumn(cs.str(), COL_WIDTHS[0])
             << PadColumn(expr_buffer.substr(idx), COL_WIDTHS[1])
             << PadColumn(_formula, COL_WIDTHS[2]) << "\n";
    };

    print_process(-1);

    while (idx < expr_buffer.size())
    {
        char cur = expr_buffer[idx];
        char top_stack = stk.back();

        if (top_stack == '#' && cur == '#')
        {
            success = 1; break;
        }

        if (!syntax.vts.count(cur) && cur != '#')
        {
            std::stringstream msg;
            msg << "Unrecognized symbol of '" << cur << "'.";
            throw std::runtime_error(msg.str());
        }

        if (syntax.vts.count(top_stack))
        {
            // matching elimination of terminal symbols
            if (cur == top_stack)
            {
                stk.pop_back();  ++idx;  ++step;
                print_process(-1);
                continue;
            }
            else
            {
                std::stringstream msg;
                msg << "Terminal symbol not matching: '" << top_stack << "' and '" << cur << "'.";
                throw std::runtime_error(msg.str());
            }
        }

        auto key = common::CharsToUint16(top_stack, cur);
        if (!table.count(key))
        {
            std::stringstream msg;
            msg << "No matching production rules of [" << top_stack << " -> " << cur << "].";
            throw std::runtime_error(msg.str());
        }
        auto production_rule = syntax.rules[table[key]];
        stk.pop_back();
        auto formula = production_rule.second;
        // Reverse push the right part of the production rule into the stack.
        for (int i = formula.size() - 1; i >= 0; --i)
        {
            if (formula[i] != syntax.epsilon)
                stk.push_back(formula[i]);
        }
        ++step;
        print_process(table[key]);
        // Optional. Add a gap to watch the process.
        sleep(1);
    }

    std::cout << std::flush;
    std::cout << "\n";
    
    if (success)
    {
        std::cout << "Analyzing successfully.\n";
    } 
    else
    {
        std::cout << "Unknown Error Occurred.\n";
    }
}

    
} // namespace parser